# 基于matplotlib 使用 yfinance ohlc数据 实现中性杠杆网格回测

import yfinance as yf
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.dates import DateFormatter


class NeutralGridStrategy:
    def __init__(
        self,
        ticker,
        initial_capital=100000,
        grid_range=0.2,
        grid_num=10,
        leverage=3,
        fee=0.0005,
        start="2020-01-01",
        end="2025-02-23",
    ):
        self.ticker = ticker
        self.initial_capital = initial_capital
        self.grid_range = grid_range  # 价格波动范围
        self.grid_num = grid_num  # 网格层数
        self.leverage = leverage  # 单层杠杆倍数
        self.fee = fee  # 交易费率
        self.start = start
        self.end = end

        self.data = self._download_data()
        self._generate_grid()
        self._run_backtest()

    def _download_data(self):
        df = yf.download(self.ticker, start=self.start, end=self.end)
        return df[["Open", "High", "Low", "Close"]]

    def _generate_grid(self):
        """动态生成网格体系"""
        base_price = self.data["Close"].iloc
        self.grid_levels = np.linspace(
            base_price * (1 - self.grid_range),
            base_price * (1 + self.grid_range),
            self.grid_num,
        )
        self.grid_spacing = self.grid_levels - self.grid_levels

    def _run_backtest(self):
        """中性策略核心逻辑"""
        capital = self.initial_capital
        position = {"long": 0, "short": 0}  # 多空对冲仓位
        equity = [capital]
        trade_log = []

        for close_price in self.data["Close"]:
            # 计算价格偏离度
            deviation = (close_price - self.grid_levels.mean()) / self.grid_spacing

            # 杠杆平衡公式
            target_long = max(0, deviation) * self.leverage
            target_short = max(0, -deviation) * self.leverage

            # 执行调仓
            delta_long = target_long - position["long"]
            delta_short = target_short - position["short"]

            # 计算交易成本
            fee = abs(delta_long + delta_short) * close_price * self.fee
            capital -= fee

            # 更新仓位和资金
            position = {"long": target_long, "short": target_short}
            equity.append(
                capital + (position["long"] - position["short"]) * close_price
            )

            trade_log.append(
                {
                    "price": close_price,
                    "long": target_long,
                    "short": target_short,
                    "fee": fee,
                }
            )

        self.equity = equity
        self.trade_log = pd.DataFrame(trade_log)

    def visualize(self):
        """可视化模块"""
        plt.figure(figsize=(14, 8))

        # 价格与网格体系
        ax1 = plt.subplot2grid((4, 1), (0, 0), rowspan=3)
        ax1.plot(self.data["Close"], label="Price", linewidth=1)
        for level in self.grid_levels:
            ax1.axhline(level, color="gray", linestyle="--", alpha=0.5)
        ax1.set_title(f"Neutral Grid Strategy - {self.ticker}")

        # 多空仓位
        ax2 = plt.subplot2grid((4, 1), (3, 0), sharex=ax1)
        ax2.plot(self.trade_log["long"], label="Long Position", color="green")
        ax2.plot(self.trade_log["short"], label="Short Position", color="red")
        ax2.set_ylabel("Leverage")

        # 格式化显示
        ax1.xaxis.set_major_formatter(DateFormatter("%Y-%m"))
        plt.gcf().autofmt_xdate()
        ax1.legend()
        ax2.legend()
        plt.tight_layout()

        # 资金曲线
        plt.figure(figsize=(14, 4))
        plt.plot(self.equity, label="Equity Curve")
        plt.title(
            f"Total Return: {((self.equity[-1] / self.initial_capital - 1) * 100):.2f}%"
        )
        plt.xlabel("Trading Days")
        plt.show()


# 策略实例化与运行
strategy = NeutralGridStrategy(
    ticker="BTC-USD",
    initial_capital=100000,
    grid_range=0.3,
    grid_num=20,
    leverage=5,
    start="2024-02-15",
)
strategy.visualize()
